import cv2
import sys
import SVM_ocr
import numpy as np
import matplotlib.pyplot as plt


def display(window, img, width, block=False):
    """显示图片"""
    cv2.namedWindow(window, cv2.WINDOW_NORMAL)
    cv2.resizeWindow(window, width, int(width * img.shape[0] / img.shape[1]))
    cv2.imshow(window, img)
    if block:
        cv2.waitKey(0)
        cv2.destroyAllWindows()


def close_op(img, size):
    """形态学闭合操作"""
    pre_element = cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2))
    ero_element = cv2.getStructuringElement(cv2.MORPH_RECT, (size, int(size / 8)))
    dil_element = cv2.getStructuringElement(cv2.MORPH_RECT, (size, int(size / 6)))
    pre = cv2.erode(img, pre_element, iterations=1)
    dil = cv2.dilate(pre, dil_element, iterations=1)
    ero = cv2.erode(dil, ero_element, iterations=1)
    return ero


def close_op_y(img, size):
    ero_element = cv2.getStructuringElement(cv2.MORPH_RECT, (int(size / 8), size))
    dil_element = cv2.getStructuringElement(cv2.MORPH_RECT, (int(size / 6), size))
    dil = cv2.dilate(img, dil_element, iterations=1)
    ero = cv2.erode(dil, ero_element, iterations=1)
    return ero


def check_plate(img, rect):
    """检查车牌是否存在"""
    return True


def draw_rect(img, rect):
    """在图像中绘制矩形区域"""
    box = cv2.boxPoints(rect)
    box = np.int0(box)
    return cv2.polylines(img, [box], isClosed=True, color=(0, 0, 255), thickness=1)


def scale_rect(rect, scale=1.0):
    """缩放矩形区域"""
    rect = list(rect)
    scale = scale * rect[1][0] - rect[1][0]
    rect[1] = [x + scale for x in rect[1]]
    return rect


def getBlueAndGreenChannel(img):
    """获得图像的蓝色和绿色通道"""
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    lower_blue = np.array([100, 90, 40])
    upper_blue = np.array([124, 255, 255])
    lower_green = np.array([40, 40, 40])
    upper_green = np.array([80, 255, 255])
    mask_blue = cv2.inRange(hsv, lower_blue, upper_blue)
    mask_green = cv2.inRange(hsv, lower_green, upper_green)
    mask = cv2.bitwise_or(mask_blue, mask_green)
    _, s, _ = cv2.split(hsv)
    img = cv2.bitwise_and(s, s, mask=mask)
    display('getBlueAndGreenChannel', img, 300, block=False)
    return img


def adjust_box(img, box):
    """调整车牌区域"""
    pts = np.array(np.where(img > 0)).T
    pts[:, [0, 1]] = pts[:, [1, 0]]
    new_box = []
    for corner in box:
        new_box.append(list(pts[np.argmin([np.linalg.norm(pt - corner) for pt in pts])]))
    new_box = np.array(new_box)
    new_box[0][0] -= 3
    new_box[3][0] -= 3
    return new_box


def box_normalize(box):
    box = box[np.argsort(box[:, 1])]
    if box[0][0] > box[1][0]:
        box[[0, 1], :] = box[[1, 0], :]
    if box[2][0] < box[3][0]:
        box[[2, 3], :] = box[[3, 2], :]
    return box


def extract_plate(src, rect):
    """从矩形区域中提取车牌"""
    rect = scale_rect(rect, 1.2)

    """裁剪矩形区域"""
    img = np.zeros_like(src).astype('uint8')
    box = np.int0(cv2.boxPoints(rect))
    box = box_normalize(box)
    img = cv2.fillPoly(img, [box], (255, 255, 255))
    cv2.copyTo(src, img, img)

    """图像预处理操作"""
    img = getBlueAndGreenChannel(img)  # 获取蓝色和绿色通道
    img = cv2.GaussianBlur(img, (3, 3), 0, 0, cv2.BORDER_DEFAULT)
    th = (np.mean(img) + (np.max(img) - np.mean(img)) * 0.4)
    _, img = cv2.threshold(img, th, 255, cv2.THRESH_BINARY)
    img = close_op(img, 32)

    """调整车牌区域"""
    image = cv2.polylines(src, [box], isClosed=True, color=(0, 0, 255), thickness=1)
    display('image', image, 360, False)
    box = adjust_box(img, box)

    """通过透视变换纠正失真"""
    box = box.astype('float32')
    M = cv2.getPerspectiveTransform(box, np.float32([[0, 0], [179, 0], [179, 59], [0, 59]]))
    plate = cv2.warpPerspective(src, M, (180, 60))

    return plate


def plot_data(data):
    x = np.linspace(0, len(data), len(data))
    plt.plot(x, data)
    plt.show()


def calc_potential(img):
    """计算势能"""
    sum = np.sum(img, 0)
    print(sum)
    pot = np.zeros_like(sum)
    p, rp = 0, 0
    for i in range(len(sum)):
        """线性增长指数衰减"""
        p = p + sum[i] if sum[i] else p / 5
        rp = rp + sum[len(sum) - i - 1] if sum[len(sum) - i - 1] else rp / 5
        pot[i] += p
        pot[len(sum) - i - 1] += rp
    return pot


def adjust_vision(img):
    """调整视场大小以更好地匹配"""
    sum = np.sum(img, 1)
    up, down = -1, -1
    1
    for i in range(len(sum)):
        if sum[i] > 512:
            if up == -1:
                up = i
            down = i
    height = down - up + 1
    img = img[up: down + 1, :]
    padding = height - img.shape[1]
    left = (int)(padding / 2)
    right = (int)(padding - left)
    img = np.hstack((np.zeros((height, left)), img))
    img = np.hstack((img, np.zeros((height, right))))
    img = cv2.resize(img, (20, 20))
    return img


def plate_split(img):
    """通过势能拆分字符"""
    pot = calc_potential(img)
    limit_pot = np.mean(pot) * 0.5
    print('limit potential:', limit_pot)
    pot[pot < limit_pot] = 0

    mid_points = []
    peak_points = []

    last = pot[0]
    for i in range(1, len(pot) - 1):
        if pot[i] > pot[i - 1] and pot[i] > pot[i + 1]:
            peak_points.append(i)
        elif pot[i] < pot[i - 1] and pot[i] < pot[i + 1]:
            mid_points.append(i)

    if len(mid_points) == 0:
        mid_points = [i for i in range(1, len(pot) - 1)]

    gap = []
    for i in range(len(mid_points) - 1):
        gap.append(mid_points[i + 1] - mid_points[i])

    avg_gap = np.mean(gap) if len(gap) > 0 else 0

    candidates = []
    for mid in mid_points:
        if pot[mid] > 0 and (len(candidates) == 0 or mid - candidates[-1] > avg_gap * 0.5):
            candidates.append(mid)

    chars = []
    start = 0
    for mid in candidates:
        chars.append(img[:, start:mid])
        start = mid
    chars.append(img[:, start:])

    chars = [adjust_vision(char) for char in chars]
    return chars


def plate_cut_text(img):
    """车牌字符分割，将分割好的车牌区域进行字符分割"""
    sum = np.sum(img, 1)
    limit = np.mean(sum) * 0.2
    bound = []
    start = 0
    for i in range(len(sum) - 1):
        if sum[i] < limit and sum[i + 1] >= limit:
            start = i
        elif sum[i] >= limit and sum[i + 1] < limit:
            bound.append((start, i))

    chars = []
    for b in bound:
        if b[1] - b[0] > 2:  # 忽略过小的分割区域
            char = img[b[0]:b[1], :]
            char = adjust_vision(char)
            chars.append(char)
    return chars


def recognize_plate(img, reader):
    """识别车牌字符"""
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    _, img = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    img = close_op_y(img, 10)

    chars = plate_cut_text(img)
    if len(chars) != 7:
        print(f"识别到的字符数量不对: {len(chars)}")
        return ""

    results = []
    for i, char in enumerate(chars):
        if i == 0:
            result = reader.recognize_chinese(char)
        else:
            result = reader.recognize_alnum(char)
        results.append(result[0])

    return "".join(results)


def main(image_path):
    reader = SVM_ocr.Reader()

    image = cv2.imread(image_path)
    display('Original Image', image, 800, False)

    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    blur = cv2.GaussianBlur(gray, (5, 5), 0)
    edges = cv2.Canny(blur, 100, 200)

    contours, _ = cv2.findContours(edges, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    contours = sorted(contours, key=cv2.contourArea, reverse=True)[:10]

    for contour in contours:
        rect = cv2.minAreaRect(contour)
        if check_plate(gray, rect):
            plate = extract_plate(image, rect)
            display('Extracted Plate', plate, 800, False)
            text = recognize_plate(plate, reader)
            print(f"识别出的车牌: {text}")
            return text


if __name__ == '__main__':
    if len(sys.argv) != 2:
        print("Usage: python main.py <path_to_image>")
        sys.exit(1)
    image_path = sys.argv[1]
    main(image_path)
